import random
import gym
import torch
import torch.nn.functional as F
import numpy as np
import utils
from Net import PolicyNet, ValueNet, PolicyNetContinuous


class PPO:
    """PPO-clip算法"""

    def __init__(self, state_dim, hidden_dim, action_dim, actor_lr, critic_lr, lmbda, epochs, eps, gamma, device,
                 action_type="discrete"):
        if action_type == "continuous":
            self.actor = PolicyNetContinuous(state_dim, hidden_dim, action_dim).to(device)
        else:
            self.actor = PolicyNet(state_dim, hidden_dim, action_dim).to(device)
        self.critic = ValueNet(state_dim, hidden_dim).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_lr)
        self.gamma = gamma
        self.lmbda = lmbda
        # 一条序列的数据用来训练轮数
        self.epochs = epochs
        # PPO-clip中截断范围的参数
        self.eps = eps
        self.device = device
        self.action_type = action_type

    def take_action(self, state):
        state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
        if self.action_type == "continuous":
            mu, std = self.actor(state)
            action_list = torch.distributions.Normal(mu, std)
            return [action_list.sample().item()]
        else:
            probs = self.actor(state)
            action_list = torch.distributions.Categorical(probs)
            action = action_list.sample()
            return action.item()

    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)
        if self.action_type == "continuous":
            actions = torch.tensor(transition_dict['actions'], dtype=torch.float).view(-1, 1).to(self.device)
        else:
            actions = torch.tensor(transition_dict['actions'], dtype=torch.int64).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'], dtype=torch.float).view(-1, 1).to(self.device)

        # 修改奖励(用于Pendulum)
        # rewards = (rewards + 8.0) / 8.0

        td_target = rewards + self.gamma * self.critic(next_states) * (1 - dones)
        td_error = td_target - self.critic(states)
        # 计算优势
        advantage = utils.compute_advantage(self.gamma, self.lmbda, td_error.cpu()).to(self.device)
        if self.action_type == "continuous":
            mu, std = self.actor(states)
            action_dists = torch.distributions.Normal(mu.detach(), std.detach())
            # 动作是正态分布
            old_log_probs = action_dists.log_prob(actions)
        else:
            old_log_probs = torch.log(self.actor(states).gather(1, actions)).detach()

        for _ in range(self.epochs):
            if self.action_type == "continuous":
                mu, std = self.actor(states)
                action_dists = torch.distributions.Normal(mu, std)
                # 动作是正态分布
                log_probs = action_dists.log_prob(actions)
            else:
                log_probs = torch.log(self.actor(states).gather(1, actions))

            # import sampling
            ratio = torch.exp(log_probs - old_log_probs)
            surr1 = ratio * advantage
            surr2 = torch.clamp(ratio, 1 - self.eps, 1 + self.eps) * advantage  # 截断
            # PPO损失函数
            actor_loss = torch.mean(-torch.min(surr1, surr2))
            # 价值网络损失函数
            critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))
            self.actor_optimizer.zero_grad()
            self.critic_optimizer.zero_grad()
            actor_loss.backward()
            critic_loss.backward()
            self.actor_optimizer.step()
            self.critic_optimizer.step()


if __name__ == '__main__':
    actor_lr = 1e-4
    critic_lr = 5e-3
    num_episodes = 2000
    hidden_dim = 128
    gamma = 0.9
    lmbda = 0.9
    epochs = 10
    epsilon = 0.2
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device(
        "cpu")

    env_name = 'Pendulum-v1'
    env = gym.make(env_name)
    random.seed(0)
    np.random.seed(0)
    env.seed(0)
    torch.manual_seed(0)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    agent = PPO(state_dim, hidden_dim, action_dim, actor_lr, critic_lr, lmbda, epochs, epsilon, gamma, device,
                action_type="continuous")

    return_list = utils.train_on_policy_agent(env, agent, num_episodes)
    utils.plots(return_list, env_name, "PPO-clip")
